import numpy as np
from keras.models import Sequential
from keras.layers.core import Dense, Activation, Dropout
from keras.layers.recurrent import LSTM
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt

#固定的训练长度
fix_train_delta = 4
fix_train_min = fix_train_delta + 1
fix_train_max = 50

def predict(x_train, y_train, x_test):
    # y_train = np.array(y)
    # x_train = np.reshape( x, (x.shape[0], x.shape[1], 1) )#Lstm调用库函数必须要进行维度转换
    model = Sequential()
    model.add( LSTM( 50, input_shape=(x_train.shape[1], x_train.shape[2]), return_sequences=True) )
    model.add( LSTM( 50, return_sequences=False ) )
    # model.add( Dropout( 0.2 ) )
    model.add( Dense( 1 ) )
    model.add( Activation( 'linear' ) )
    model.compile( loss="mean_squared_error", optimizer="adam" )
    model.fit( x_train, y_train, epochs=3, batch_size=1, verbose=2)#参数依次为特征，标签，训练循环次数，小批量（一次放入训练的数据个数）
    res = model.predict(x_test)
    return res

def main():
    df = pd.read_csv('../stick-data-day/sh600000.csv')
    new_data = pd.DataFrame(index=range(0,len(df)),columns=['close'])
    for i in range(0, len(df)):
        new_data['close'][i] = df['close'][i]
    dataset = new_data.values
    if (len(dataset) < fix_train_min):
        print("data is too little")
        return
    if (len(dataset) > fix_train_max):
        dataset = dataset[(len(dataset) - fix_train_max):]
    #converting dataset into x_train and y_train
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(dataset)
    x_train, y_train, x_test = [],[],[]
    for i in range(fix_train_delta, len(scaled_data)):
        x_train.append(scaled_data[i-fix_train_delta:i, 0])
        y_train.append(scaled_data[i, 0])
    x_test.append(scaled_data[len(scaled_data) - fix_train_delta:])
    x_train, y_train, x_test = np.array(x_train), np.array(y_train), np.array(x_test)
    x_train = x_train.reshape(x_train.shape[0], x_train.shape[1], 1)
    x_test = x_test.reshape(x_test.shape[0], x_test.shape[1], 1)
    res = predict(x_train, y_train, x_test)
    res = scaler.inverse_transform(res)
    predict_len = 10
    predict_data = pd.DataFrame(index=range(0, predict_len),columns=['close'])
    for i in range(0, len(predict_data)):
        predict_data["close"][i] = res[0]
    predict_data.index = range(len(new_data),len(new_data) + predict_len)
    print(predict_data)
    # print("new_data shape", new_data.shape)
    # print("predict_data shape", predict_data.shape)
    plt.plot(new_data["close"])
    plt.plot(predict_data["close"])
    plt.show()

main()